import pandas as pd
import numpy as np
from itertools import combinations
import xlrd

def preprocess_dataframe(df):
    
    # Extract 
    df[['Concen', '极性']] = df['文件'].str.extract(r'CBY-([^-]+)-([^.]+)\.d')
    
    # Create m/z-质量
    df['差'] = df['m/z']-df['质量']
    
    return df

def transform_dataframe(df):

    # Create new dataframe with specified columns
    new_df = pd.DataFrame()
    
    # 1. Empty "No." column
    new_df['No.'] = ''
    
    # 2-7. Copy specified columns
    new_df['Concen'] = df['Concen']
    new_df['Mode'] = df['极性']
    new_df['Name'] = df['名称']
    new_df['Formula'] = df['分子式'].str.replace(' ', '')
    new_df['RT'] = df['RT']
    new_df['m/z'] = df['m/z']
    new_df['Mass'] = df['质量']
    
    # 8. Create Adduct column based on "差" values
    def determine_adduct(x):
        if 0.95 <= x <= 1.05:
            return '[M+H]+'
        elif -1.05 <= x <= -0.95:
            return '[M-H]-'
        elif 44.95 <= x <= 45.05:
            return '[M+COOH]-'
        elif 22.95 <= x <= 23.05:
            return '[M+Na]+'
        elif 17.95 <= x <= 18.05:
            return '[M+NH4]+'
        return ''
    
    new_df['Adduct'] = df['差'].apply(determine_adduct)
    
    # 9-12. Copy specified columns
    new_df['Mass error'] = df['误差(Tgt, ppm)']
    new_df['Score'] = df['分数 (Tgt)']
    new_df['Height'] = df['峰高']
    new_df['Area'] = df['面积']
    
    # 13-15. Empty columns
    new_df['峰面积百分比'] = ''
    new_df['CID'] = ''
    new_df['CAS'] = ''
    new_df['SMILES'] = ''
    
     # 16. Create "正负对应" column
    new_df['正负对应'] = ''
    
    # Round Mass to 2 decimals for comparison
    df['Mass_rounded'] = df['质量'].round(2)
    
    # Create groups for "正负对应"
    pos_neg_group = 1
    processed_indices = set()
    
    # First, create a dictionary to store all potential matches
    potential_matches = {}
    
    # Iterate through all pairs of rows
    for i in range(len(df)):
        for j in range(i + 1, len(df)):
            # Check if rows match the criteria
            if (df.iloc[i]['Concen'] == df.iloc[j]['Concen'] and
                df.iloc[i]['Mass_rounded'] == df.iloc[j]['Mass_rounded'] and
                df.iloc[i]['极性'] != df.iloc[j]['极性'] and
                abs(df.iloc[i]['RT'] - df.iloc[j]['RT']) < 0.5):
                
                # Store both indices as a frozen set (for uniqueness)
                match_pair = frozenset([i, j])
                potential_matches[match_pair] = True
    
    # Now process all matches to create groups
    while potential_matches:
        # Start with first pair
        first_pair = list(potential_matches.keys())[0]
        current_group = set(first_pair)
        
        # Keep track of pairs to remove
        pairs_to_remove = {first_pair}
        
        # Keep expanding the group until no new matches are found
        changed = True
        while changed:
            changed = False
            for pair in potential_matches.keys():
                if pair not in pairs_to_remove:
                    # If this pair shares any index with current group
                    if any(idx in current_group for idx in pair):
                        current_group.update(pair)
                        pairs_to_remove.add(pair)
                        changed = True
        
        # Assign group number to all indices in the group
        for idx in current_group:
            new_df.loc[idx, '正负对应'] = str(pos_neg_group)
            processed_indices.add(idx)
        
        # Remove processed pairs
        for pair in pairs_to_remove:
            del potential_matches[pair]
        
        pos_neg_group += 1
    
    # 17. Create "同时出峰" column
    new_df['同时出峰'] = ''
    
    # Reset processed indices for new grouping
    processed_indices = set()
    same_peak_group = 1
    
    for i in range(len(df)):
        if i in processed_indices:
            continue
            
        current_polarity = df.iloc[i]['极性']
        current_mass = df.iloc[i]['Mass_rounded']
        current_rt = df.iloc[i]['RT']
        current_buwei = df.iloc[i]['Concen']
        
        # Find matching rows
        matching_rows = []
        for j in range(len(df)):
            if i != j and j not in processed_indices:
                if (df.iloc[j]['极性'] == current_polarity and
                    df.iloc[j]['Mass_rounded'] == current_mass and
                    df.iloc[j]['Concen'] != current_buwei and
                    abs(df.iloc[j]['RT'] - current_rt) < 0.5):
                    matching_rows.append(j)
        
        if matching_rows:
            matching_rows.append(i)
            for idx in matching_rows:
                new_df.loc[idx, '同时出峰'] = str(same_peak_group)
                processed_indices.add(idx)
            same_peak_group += 1
    
    new_df['tgt'] = df['标记 (Tgt)'].str.replace(' ', '')
    
    return new_df


def filter_low_scores(df):
    
    # Create a copy
    filtered_df = df.copy()
    
    # Filter satisfied conditions
    filtered_df = filtered_df[
        ((filtered_df['正负对应'] != '') |
         (filtered_df['同时出峰'] != '') |
         (filtered_df['Score'] >= 75))
    ] 
    
    filtered_df = filtered_df.reset_index(drop=True)
    
    return filtered_df


def round_decimals(df):
    
    # Round decimals
    df['m/z'] = df['m/z'].round(2)
    df['Mass'] = df['Mass'].round(2)
    
    return df


df = pd.read_excel('CBY-初始数据.xlsx') #, engine='xlrd')

processed_df = preprocess_dataframe(df)

transformed_df = transform_dataframe(processed_df)

filtered_df = filter_low_scores(transformed_df)

percent_df = filtered_df.copy()                    
area_sums = percent_df.groupby(['Concen', 'Mode'])['Area'].transform('sum')
percent_df['峰面积百分比'] = (percent_df['Area'] / area_sums * 100).round(2)
# percent_df['m/z'] = percent_df['m/z'].round(2)
# percent_df['Mass'] = percent_df['Mass'].round(2)
 
percent_df.to_excel('processed_file.xlsx', index=False) 